import requests
import difflib
import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
import os
from math import floor
import numpy as np
from PIL import Image
import sys
from safebooru.config import device


class CHEWANbooru(Dataset):
    """CHEWAN Dataset from <https://yande.re/post/> ."""
    from safebooru.labels import main_chara

    def __init__(self, root_dir="./safebooru/", train=True, transform=None):
        """
        Args:
            root_dir (string): Path to the folder where '/label' and '/dataset' exitst.
            train (bool): If true, return first 80 percent of the data, otherwise other 20 percent.
            transform (callable, optional) Optional transform to be applied on a sample.
            allowmulti (bool): If true, data with multiple classes will be excluded.
        Dataset Structure:
            label
                labels one-hot-encoded with `numpy.save()`. Should be load with `numpy.load()`.
            dataset
                Images saved.
                cfg indicates the used images.
                new indicates the number of previously used images.
        """
        self.root = root_dir
        self.train = train
        self.transform = transform
        self.datalist = []
        self.labelist = []

        for k,i in enumerate(self.main_chara):
            lst = os.listdir(f"{root_dir}dataset224/{i}/")
            num = len(lst)
            if num < 10:
                print(f"{i} has only {num} valid datas")
            sep = floor(num*0.8)
            if train:
                z = lst[:sep]
            else:
                z = lst[sep:]
            self.datalist += [int(os.path.splitext(k)[0])for k in z]
            self.labelist += [k] * len(z)
        # print(sys.getsizeof(self.datalist)+sys.getsizeof(self.labelist))

    def __len__(self):
        return len(self.datalist)

    def __getitem__(self, idx):
        """
        Args:
            idx (int): Index

        Returns:
            tuple: (image, target) where target is the index target class.
        """
        imgpath = self.root+"dataset224/"+self.main_chara[self.labelist[idx]]+"/"+str(self.datalist[idx])+".jpg"
        label = self.labelist[idx]
        image = Image.open(imgpath)
        # print(imgpath)
        if self.transform:
            image = self.transform(image)
        # label = torch.from_numpy(np.load(labelpath))

        return image, label


def string_similar(s1, s2):
    return difflib.SequenceMatcher(None, s1, s2).quick_ratio()


class datasetmanager:
    def getcnt(self, label):
        lst = os.listdir(f"./dataset/{label}/")
        if "cfg" in lst:
            lst.remove("cfg")
        if "new" in lst:
            lst.remove("new")
        return len(lst)
    from labels import main_chara

    def getLabel(self, label):
        assert len(label) == len(self.main_chara)
        for i in self.main_chara:
            if label[self.main_chara[i]["index"]] == 1:
                return i
        assert False

    def addpic_(self, url, label):
        Label = self.getLabel(label)
        cnt = self.getcnt(Label)
        filename = os.path.basename(url)
        for k in range(20):
            r = requests.get(url, stream=True)
            if r.status_code == 200:
                with open(f"./dataset/{Label}/{cnt}{os.path.splitext(filename)[-1]}", 'wb') as f:
                    for chunk in r:
                        f.write(chunk)
                    # print(f"save pic at ./dataset/{Label}/{cnt}{os.path.splitext(filename)[-1]}")
                    break
            else:
                print(f"第{k}次針下{Label} {cnt}")
        np.save(f"./label/{Label}/{cnt}.npy", label)
        # print(f"to save npy at ./label/{Label}/{cnt}.npy")

    def getsimilar(self, inp):
        for i in self.main_chara:
            if string_similar(inp, self.main_chara[i]["chn"]) > 0.9:
                return i
            if string_similar(inp, self.main_chara[i]["eng"]) > 0.9:
                return i
            if string_similar(inp, self.main_chara[i]["jpn"]) > 0.9:
                return i
        return None

    def addpic(self):
        while True:
            url = input("url:")
            label = np.zeros(len(self.main_chara))
            k = input("label:")
            while k != "no":
                z = self.getsimilar(k)
                if z is None:
                    print("kbd")
                else:
                    print(f"{z} recorded")
                    label[self.main_chara[z]["index"]] = 1
                k = input("label:")
            self.addpic_(url, label)


if __name__ == "__main__":
     z = datasetmanager()
     z.addpic()
